#ifndef KPEXTCOMPUTEPIPELINE_H
#define KPEXTCOMPUTEPIPELINE_H

#include "kpextCommonDef.h"

namespace kpext{

//------------------- 声明 -----------------
class ComputePipeline;
class PipelineConstSZ;
//-----------------------------------------

/**
 * @brief 生成 vulkan 计算管线的简单封装类
 */
class ComputePipeline
{
public:
    explicit ComputePipeline(std::shared_ptr<vk::Device> device,
                             std::shared_ptr<vk::PipelineCache> globaPipelinecache = nullptr
                             );
    virtual ~ComputePipeline();


public:

    //----------下面的成员变量更改后调用rebuild可重新建立计算管线------
    std::vector<std::shared_ptr<vk::DescriptorSetLayout>> m_descrSetLayouts;
    std::shared_ptr<PipelineConstSZ> m_specConsts;
    std::shared_ptr<PipelineConstSZ> m_pushConsts;

    /**
     * @brief 是否使用全局管线缓存， 以加快速pipeline重新创建速度
     */
    bool m_useGlobalPipelinecache;

    //------------------------------------------------------


    /**
     * @brief 重新构建计算管线，除上述的成员变量外，成功后下面的成员变量都可能被更改
     * @param codesize  spirv二进制代码字节大小
     * @param codedata  spirv二进制代码
     * @return  成功返回true, 否则返回false
     */
    bool rebuild(size_t codesize,
                 const void * const codedata
                 );

    /**
     * @brief 重新构建计算管线的重载版本，用本类对象里现有信息重新，成功后下面的成员变量都可能被更改
     * @return   成功返回true, 否则返回false
     */
    bool rebuild();


    //----------下面的成员变量应尽量不在类外部进行变更，也不要在类外释放------
public:

    /**
     * @brief spirv二进制代码
     */
    std::vector<uint8_t> m_spirvCode;

    /**
     * @brief vulkan着色器模块
     */
    std::shared_ptr<vk::ShaderModule> m_shaderModule;

    /**
     * @brief vulkan 管线布局
     */
    std::shared_ptr<vk::PipelineLayout> m_layout;

    /**
     * @brief vulkan 管线缓存
     */
    std::shared_ptr<vk::PipelineCache> m_pipelineCache;


    /**
     * @brief vulkan 管线
     */
    std::shared_ptr<vk::Pipeline> m_pipeline;

    //-------------------------------------------------------


protected:
    /**
     * @brief vulkan 逻辑设备
     */
    std::shared_ptr<vk::Device> m_device;

    /**
     * @brief 全局管线缓存
     */
    std::shared_ptr<vk::PipelineCache> m_GPipelineCache;


    /**
     * @brief 创建vulkan 着色器模块
     * @param codesize  spirv二进制代码字节大小
     * @param codedata  spirv二进制代码
     * @param result  成功返回vulkan shadermodule对象共享指针
     * @return  成功返回true, 否则返回false
     */
    virtual bool createShaderModule(size_t codesize,
                                    const void * const codedata,
                                    std::shared_ptr<vk::ShaderModule> & result);


    /**
     * @brief 创建vulkan pipelinelayout 管线布局
     * @param descrLayouts  描述符集布局
     * @param pushConsts  推送常量
     * @param result  成功返回vulkan pipelinelayout 对象共享指针
     * @return  成功返回true, 否则返回false
     */
    virtual bool createPipelinelayout(
            std::vector<std::shared_ptr<vk::DescriptorSetLayout>> descrLayouts,
            std::shared_ptr<PipelineConstSZ> pushConsts,
            std::shared_ptr<vk::PipelineLayout> & result
            );

    virtual std::shared_ptr<vk::Pipeline> createPipeline(
            vk::PipelineLayout & pipelineLayout,
            vk::ShaderModule & shaderModule,
            std::shared_ptr<PipelineConstSZ> specConsts,
            std::shared_ptr<vk::PipelineCache> & pipelineCache
            );

}; // class kpext::ComputePipeLine


//-------------------------------------------

/**
 * @brief 管线使用的推送常量、特化常量简单封装类
 */
class PipelineConstSZ
{
public:

    /**
     * @brief 以指定元素个数和每个元素大小构造未初始化数据的常量结构类
     * @param elemssize  元素个数
     * @param elemsbyte  元素字节大小
     */
    PipelineConstSZ(uint32_t elemssize,
                          uint32_t elemsbyte = 4
            ): PipelineConstSZ(){
        if(elemsbyte % 4){
            SPDLOG_INFO(__FUNCTION__ " Please ensure alignment ok!");
        }
        if(elemssize < 1 ||elemsbyte < 1){
            SPDLOG_ERROR("Construct a null SpecializationConstSZ object!");
            return;
        }
        m_size = elemssize * elemsbyte;
        m_data = std::shared_ptr<uint8_t>(new uint8_t[m_size]);
        for (uint32_t i = 0; i < elemssize; ++i) {
            m_elemInfos.push_back({i * elemsbyte, elemsbyte});
        }
    }


    /**
     * cpp强制结构对齐的方法如下
     * struct UniformBufferObject {
     *         alignas(16) glm::mat4 model;    // 显式指定地址对齐， 减少与glsl uniform结构体数据对齐不一致的问题
     *         alignas(16) glm::mat4 view;
     *         alignas(16) glm::mat4 proj;
     *     }; //  统一缓冲对象数据结构， 常用于着色器可访问的一致模型、视图、投影常量
     *     glsl 里变量的结构基本类型对齐要求
     * 标量必须按 N 对齐（= 4 个字节，给定 32 位浮点数）。
     * vec2 必须对齐 2N（= 8 个字节）
     * vec3 或 vec4 必须对齐 4N（= 16 字节）
     * 嵌套结构必须通过其成员的基本对齐方式进行对齐，四舍五入到 16 的倍数。
     * mat4 矩阵必须与 vec4 具有相同的对齐方式。
     * 更详细的标准 https://registry.khronos.org/vulkan/specs/1.1-extensions/html/chap15.html#interfaces-resources-layout
    */

    /**
     * @brief 从外部结构体构造常量结构类
     * @param bytesize  外部结构体字节大小
     * @param start  外部结构体指针， 不能为nullptr
     * @param elems  外部结构体各个成员的指针，请依序放入结构中各成员的指针
     * @param elemsbyte  每个成员占用的字节大小, 请依序放入结构中各成员的类型大小
     */
    PipelineConstSZ(size_t bytesize,
                          const void * const start,
                          const std::vector<void *> & elems,
                          const std::vector<uint32_t> & elemsbyte
                          ) : PipelineConstSZ()
    {
        if(bytesize < 1
                || !start
                || elems.size() < 1
                || elemsbyte.size() != elems.size()
                ){
            SPDLOG_ERROR("Construct a null SpecializationConstSZ object!");
            return;
        }

        bool check = true;
        auto sptr = reinterpret_cast<const uint8_t * const>(start);
        auto eptr = sptr + bytesize;
        size_t count = elems.size();
        for (size_t i = 0; i < count; ++i) {
            auto ptr = reinterpret_cast<decltype (sptr)>(elems[i]);
            if(ptr < sptr
                    || ptr >= eptr
                    || elemsbyte[i] < 1){
                check = false;
                break;
            }
            auto offset = static_cast<uint32_t>(ptr - sptr);
            m_elemInfos.push_back({offset, elemsbyte[i]});
        }

        if(!check){
            m_elemInfos.clear();
            SPDLOG_ERROR("Construct a null SpecializationConstSZ object!"
                         "(invalid input)");
            return;
        }

        m_size = bytesize;
        m_data = std::shared_ptr<uint8_t>(new uint8_t[m_size]);
        std::memcpy(m_data.get(), start, m_size);
    }

    virtual ~PipelineConstSZ(){}

    /**
     * @brief 以vector 数据初始化构造常量结构类
     * @param data  初始化数据
     */
    template<typename T>
    explicit PipelineConstSZ(const std::vector<T> & data
                             ) : PipelineConstSZ()
    {
        if(data.size() < 1){
            SPDLOG_ERROR("Construct a null SpecializationConstSZ object!");
            return;
        }
        size_t itemcount = data.size();
        uint32_t typebytes = static_cast<uint32_t>(sizeof (T));
        m_size = itemcount * typebytes;
        m_data = std::shared_ptr<uint8_t>(new uint8_t[m_size]);
        std::memcpy(m_data.get(), data.data(), m_size);
        for (size_t i = 0; i < itemcount; ++i) {
            auto offset = static_cast<uint32_t>(i * typebytes);
            m_elemInfos.push_back({offset, typebytes});
        }
    }

    template<typename T>
    explicit PipelineConstSZ(const std::initializer_list<T> & data) :
        PipelineConstSZ(static_cast<const std::vector<T>>(data)){}

    /**
     * @brief 把向量内容赋值给管线常量模板函数
     * @param v 数据源
     */
    template<typename T = uint32_t>
    void set(std::vector<T> const & v){
        auto count = m_size / sizeof (T);
        if(v.size() < count) count = v.size();
        auto srcptr = v.begin();
        auto dstptr = reinterpret_cast<T *>(this->data());
        for (size_t i = 0; i < count; ++i) {
            dstptr[i] = srcptr[i];
        }
    }

    /**
     * @brief 把向量内容赋值给管线常量模板函数
     * @param v 数据源
     */
    template<typename T = uint32_t>
    inline void set(std::initializer_list<T> const & v){
        set(static_cast<const std::vector<T>>(v));
    }

    /**
     * @brief 获取vector方式访问类中的数据，返回的vector是幅本不能直接访问内部数据
     * @return  std::vector
     */
    template<typename T>
    std::vector<T> vector()
    {
        auto sp = reinterpret_cast<T *>(m_data.get());
        auto ep = reinterpret_cast<T *>(m_data.get() + m_size);
        return { sp, ep };
    }

    /**
     * @brief 获取推送常量范围信息, (不会填写着色器阶段)
     * @param result  常量范围信息写入的目标结果
     */
    virtual void getPushConstantRange(
            vk::PushConstantRange & result){
            result.setOffset(0);
            result.setSize(static_cast<uint32_t>(m_size));
    }

    /**
     * @brief 获取特化常量布局信息数组
     * @param result 返回获取到的特化常量布局信息数组结果
     */
    virtual void getSpecConstEntries(
            std::vector<vk::SpecializationMapEntry> & result){
        result.clear();
        if(m_elemInfos.size() < 1) return;
        uint32_t id = 0;
        for(const auto & eleminfo : m_elemInfos){
            vk::SpecializationMapEntry entry(
                        id,
                        eleminfo.first,
                        eleminfo.second
                        );
            result.push_back(entry);
            id++;
        }
    }

    /**
     * @brief 以模板类型访问指定的元素索引
     * @param idx 元素索引，当为负时表示逆向访问， 超过范围会throw异常
     * @return
     */
    template<typename T>
    T & item(int idx){
        if(idx < 0){
            idx = static_cast<int>(m_elemInfos.size()) + idx;
        }
        if(idx >= static_cast<int>(m_elemInfos.size())){
            throw std::runtime_error(__FUNCTION__ " index invalid!");
        }

        auto offset = m_elemInfos[static_cast<size_t>(idx)].first;
        return *reinterpret_cast<T *>(m_data.get() + offset);
    }

    inline float & itemFloat(int idx){
        return item<float>(idx);
    }

    inline double & itemDouble(int idx){
        return item<double>(idx);
    }

    inline uint32_t & itemUint(int idx){
        return item<uint32_t>(idx);
    }

    inline int32_t & itemInt(int idx){
        return item<int32_t>(idx);
    }

    /**
     * @brief 以模板类型访问数据值
     * @param offset 字节偏移
     * @return
     */
    template<typename T>
    T & data(uint32_t offset){
        return *reinterpret_cast<T *>(m_data.get() + offset);
    }

    /**
     * @brief 获取内部数据指针
     * @return
     */
    virtual void * data(){return m_data.get();}

    /**
     * @brief 返回内部数据字节大小
     * @return
     */
    virtual size_t size(){return m_size;}

protected:
    /**
     * @brief 分配的内存指针
     */
    std::shared_ptr<uint8_t> m_data;

    /**
     * @brief 分配的内存字节大小
     */
    size_t m_size;

    /**
     * @brief 每个元素的信息， std::pair<uint32_t offset, uint32_t bytesize>
     */
    std::vector<std::pair<uint32_t, uint32_t>> m_elemInfos;


private:
    PipelineConstSZ() : m_data(nullptr),
        m_size(0),
        m_elemInfos()
    {
    }
}; // class kpext::SpecializationConstSZ;

} // namespace kpext



#endif // KPEXTCOMPUTEPIPELINE_H
